import os
import h5py  # 用于读写HDF5格式文件（高效存储大型多模态数据）
import torch
from torch.utils.data import Dataset, DataLoader  # PyTorch数据集基类和数据加载器


class AVEDataset(Dataset):
    """
    改进的AVE数据集类：继承PyTorch的Dataset基类，支持加载带噪声样本（纯背景数据）的音频-视频多模态数据集
    区分训练/测试阶段，训练时融合干净样本和噪声样本，测试时仅使用干净样本及细粒度标签
    """
    def __init__(self, data_root, split='train'):
        """
        初始化数据集
        Args:
            data_root: 数据集根目录路径（包含所有HDF5格式的特征、标签文件）
            split: 数据拆分类型（'train'/'test'，默认'train'），用于区分训练和测试数据加载逻辑
        """
        super(AVEDataset, self).__init__()
        self.split = split  # 数据拆分类型（训练/测试）

        # 定义核心数据文件路径（干净样本）
        self.visual_feature_path = os.path.join(data_root, 'visual_feature.h5')  # 干净视觉特征
        self.audio_feature_path = os.path.join(data_root, 'audio_feature.h5')    # 干净音频特征
        self.labels_path = os.path.join(data_root, 'labels.h5')                  # 原始细粒度标签（用于测试）
        self.dir_labels_path = os.path.join(data_root, 'mil_labels.h5')          # 视频级标签（用于训练的干净样本）
        self.sample_order_path = os.path.join(data_root, f'{split}_order.h5')    # 当前拆分的样本索引文件

        # 定义噪声样本（纯背景数据）文件路径
        self.noisy_visual_feature_path = os.path.join(data_root, 'visual_feature_noisy.h5')  # 噪声视觉特征
        self.noisy_audio_feature_path = os.path.join(data_root, 'audio_feature_noisy.h5')    # 噪声音频特征
        self.dir_labels_bg_path = os.path.join(data_root, 'labels_noisy.h5')                  # 噪声样本标签（背景标签）

        self.h5_isOpen = False  # HDF5文件是否已打开的标志（延迟加载，节省内存）

    def __getitem__(self, index):
        """
        按索引获取单一样本（支持干净样本和噪声样本）
        Args:
            index: 样本在数据集中的索引（0-based）
        Returns:
            visual_feat: 视觉特征（numpy数组）
            audio_feat: 音频特征（numpy数组）
            label: 样本标签（numpy数组，训练为视频级标签，测试为细粒度标签）
        """
        # 延迟打开HDF5文件：首次访问时打开，避免初始化时加载过多数据
        if not self.h5_isOpen:
            # 加载样本索引（干净样本的顺序）
            self.sample_order = h5py.File(self.sample_order_path, 'r')['order']
            # 加载干净样本的特征和标签
            self.visual_feature = h5py.File(self.visual_feature_path, 'r')['avadataset']
            self.audio_feature = h5py.File(self.audio_feature_path, 'r')['avadataset']
            self.clean_labels = h5py.File(self.dir_labels_path, 'r')['avadataset']

            # 训练阶段额外加载噪声样本（纯背景数据）
            if self.split == 'train':
                self.negative_labels = h5py.File(self.dir_labels_bg_path, 'r')['avadataset']  # 噪声样本标签
                self.negative_visual_feature = h5py.File(self.noisy_visual_feature_path, 'r')['avadataset']  # 噪声视觉特征
                self.negative_audio_feature = h5py.File(self.noisy_audio_feature_path, 'r')['avadataset']    # 噪声音频特征

            # 测试阶段加载细粒度标签（原始标签）
            if self.split == 'test':
                self.labels = h5py.File(self.labels_path, 'r')['avadataset']

            self.h5_isOpen = True  # 标记文件已打开

        # 计算干净样本的数量（用于区分索引对应的是干净样本还是噪声样本）
        clean_length = len(self.sample_order)
        if index >= clean_length:
            # 索引超过干净样本数量 → 加载噪声样本
            valid_index = index - clean_length  # 噪声样本在自身列表中的索引
            visual_feat = self.negative_visual_feature[valid_index]  # 噪声视觉特征
            audio_feat = self.negative_audio_feature[valid_index]    # 噪声音频特征
            label = self.negative_labels[valid_index]                 # 噪声样本标签（背景）
        else:
            # 索引在干净样本范围内 → 加载干净样本
            sample_index = self.sample_order[index]  # 映射到原始数据中的索引
            visual_feat = self.visual_feature[sample_index]  # 干净视觉特征
            audio_feat = self.audio_feature[sample_index]    # 干净音频特征
            # 训练阶段使用视频级标签，测试阶段使用细粒度标签
            if self.split == 'train':
                label = self.clean_labels[sample_index]
            else:
                label = self.labels[sample_index]

        return visual_feat, audio_feat, label

    def __len__(self):
        """
        获取数据集总样本数（训练时包含干净样本+噪声样本，测试时仅含干净样本）
        Returns:
            length: 当前拆分的总样本数量
        """
        if self.split == 'train':
            # 训练集长度 = 干净样本数 + 噪声样本数
            sample_order = h5py.File(self.sample_order_path, 'r')['order']
            noisy_labels = h5py.File(self.dir_labels_bg_path, 'r')['avadataset']
            length = len(sample_order) + len(noisy_labels)
        elif self.split == 'test':
            # 测试集长度 = 干净样本数（无噪声样本）
            sample_order = h5py.File(self.sample_order_path, 'r')['order']
            length = len(sample_order)
        else:
            # 不支持其他拆分类型（如验证集需额外实现）
            raise NotImplementedError

        return length